#!/usr/bin/env python
# -*- coding: utf-8 -*-

import csv
import json
from collections import defaultdict
import os

# 输入和输出文件
input_file = 'nobel.csv'
json_file = 'knowledge.json'

# 国家名称映射（英文到中文）
country_mapping = {
    'Afghanistan': '阿富汗',
    'Albania': '阿尔巴尼亚',
    'Algeria': '阿尔及利亚',
    'Argentina': '阿根廷',
    'Armenia': '亚美尼亚',
    'Australia': '澳大利亚',
    'Austria': '奥地利',
    'Austria-Hungary': '奥匈帝国',
    'Austrian Empire': '奥地利帝国',
    'Azerbaijan': '阿塞拜疆',
    'Belarus': '白俄罗斯',
    'Belgium': '比利时',
    'Bolivia': '玻利维亚',
    'Bosnia and Herzegovina': '波斯尼亚和黑塞哥维那',
    'Brazil': '巴西',
    'British Mandate of Palestine': '英国托管巴勒斯坦',
    'British Protectorate of Palestine': '英国保护下的巴勒斯坦',
    'Bulgaria': '保加利亚',
    'Burma': '缅甸',
    'Cambodia': '柬埔寨',
    'Cameroon': '喀麦隆',
    'Canada': '加拿大',
    'Chile': '智利',
    'China': '中国',
    'Colombia': '哥伦比亚',
    'Costa Rica': '哥斯达黎加',
    'Croatia': '克罗地亚',
    'Cuba': '古巴',
    'Cyprus': '塞浦路斯',
    'Czech Republic': '捷克共和国',
    'Czechoslovakia': '捷克斯洛伐克',
    'Democratic Republic of the Congo': '刚果民主共和国',
    'Denmark': '丹麦',
    'East Germany': '东德',
    'East Timor': '东帝汶',
    'Ecuador': '厄瓜多尔',
    'Egypt': '埃及',
    'El Salvador': '萨尔瓦多',
    'Eritrea': '厄立特里亚',
    'Estonia': '爱沙尼亚',
    'Ethiopia': '埃塞俄比亚',
    'Finland': '芬兰',
    'France': '法国',
    'French Algeria': '法属阿尔及利亚',
    'Gabon': '加蓬',
    'Georgia': '格鲁吉亚',
    'Germany': '德国',
    'Ghana': '加纳',
    'Greece': '希腊',
    'Guatemala': '危地马拉',
    'Haiti': '海地',
    'Honduras': '洪都拉斯',
    'Hungary': '匈牙利',
    'Iceland': '冰岛',
    'India': '印度',
    'Indonesia': '印度尼西亚',
    'Iran': '伊朗',
    'Iraq': '伊拉克',
    'Ireland': '爱尔兰',
    'Israel': '以色列',
    'Italy': '意大利',
    'Jamaica': '牙买加',
    'Japan': '日本',
    'Jordan': '约旦',
    'Kazakhstan': '哈萨克斯坦',
    'Kenya': '肯尼亚',
    'Korea': '朝鲜',
    'Kosovo': '科索沃',
    'Kuwait': '科威特',
    'Kyrgyzstan': '吉尔吉斯斯坦',
    'Laos': '老挝',
    'Latvia': '拉脱维亚',
    'Lebanon': '黎巴嫩',
    'Liberia': '利比里亚',
    'Libya': '利比亚',
    'Lithuania': '立陶宛',
    'Luxembourg': '卢森堡',
    'Macedonia': '马其顿',
    'Madagascar': '马达加斯加',
    'Malaysia': '马来西亚',
    'Malta': '马耳他',
    'Mauritania': '毛里塔尼亚',
    'Mauritius': '毛里求斯',
    'Mexico': '墨西哥',
    'Moldova': '摩尔多瓦',
    'Monaco': '摩纳哥',
    'Mongolia': '蒙古',
    'Montenegro': '黑山',
    'Morocco': '摩洛哥',
    'Mozambique': '莫桑比克',
    'Myanmar': '缅甸',
    'Namibia': '纳米比亚',
    'Nepal': '尼泊尔',
    'Netherlands': '荷兰',
    'the Netherlands': '荷兰',
    'New Zealand': '新西兰',
    'Nicaragua': '尼加拉瓜',
    'Niger': '尼日尔',
    'Nigeria': '尼日利亚',
    'North Korea': '朝鲜',
    'North Macedonia': '北马其顿',
    'Northern Ireland': '北爱尔兰',
    'Northern Rhodesia': '北罗得西亚',
    'Norway': '挪威',
    'Ottoman Empire': '奥斯曼帝国',
    'Pakistan': '巴基斯坦',
    'Palestine': '巴勒斯坦',
    'Panama': '巴拿马',
    'Papua New Guinea': '巴布亚新几内亚',
    'Paraguay': '巴拉圭',
    'Peru': '秘鲁',
    'Philippines': '菲律宾',
    'Poland': '波兰',
    'Portugal': '葡萄牙',
    'Prussia': '普鲁士',
    'Qatar': '卡塔尔',
    'Republic of the Congo': '刚果共和国',
    'Romania': '罗马尼亚',
    'Russia': '俄罗斯',
    'Russian Empire': '俄罗斯帝国',
    'Rwanda': '卢旺达',
    'Saudi Arabia': '沙特阿拉伯',
    'Scotland': '苏格兰',
    'Senegal': '塞内加尔',
    'Serbia': '塞尔维亚',
    'Sierra Leone': '塞拉利昂',
    'Singapore': '新加坡',
    'Slovakia': '斯洛伐克',
    'Slovenia': '斯洛文尼亚',
    'Somalia': '索马里',
    'South Africa': '南非',
    'South Korea': '韩国',
    'South Sudan': '南苏丹',
    'Southern Rhodesia': '南罗得西亚',
    'Soviet Union': '苏联',
    'Spain': '西班牙',
    'Sri Lanka': '斯里兰卡',
    'Sudan': '苏丹',
    'Sweden': '瑞典',
    'Switzerland': '瑞士',
    'Syria': '叙利亚',
    'Taiwan': '台湾',
    'Tajikistan': '塔吉克斯坦',
    'Tanzania': '坦桑尼亚',
    'Thailand': '泰国',
    'Tibet': '西藏',
    'Tunisia': '突尼斯',
    'Turkey': '土耳其',
    'Turkmenistan': '土库曼斯坦',
    'Uganda': '乌干达',
    'Ukraine': '乌克兰',
    'United Arab Emirates': '阿拉伯联合酋长国',
    'United Kingdom': '英国',
    'United States': '美国',
    'USA': '美国',
    'Uruguay': '乌拉圭',
    'USSR': '苏联',
    'Uzbekistan': '乌兹别克斯坦',
    'Vatican City': '梵蒂冈',
    'Venezuela': '委内瑞拉',
    'Vietnam': '越南',
    'Wales': '威尔士',
    'West Germany': '西德',
    'Yemen': '也门',
    'Yugoslavia': '南斯拉夫',
    'Zambia': '赞比亚',
    'Zimbabwe': '津巴布韦'
}

def convert_country(country):
    """将国家名从英文转换为中文"""
    if country and country in country_mapping:
        return country_mapping[country]
    return country

def convert_and_group_data():
    """将CSV文件中的国家名称转换为中文，按birth_country分组，并追加到knowledge.json"""
    # 读取现有的JSON文件，获取最大ID
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            existing_data = json.load(f)
            max_id = max(item['id'] for item in existing_data) if existing_data else 0
    except (FileNotFoundError, json.JSONDecodeError):
        # 如果文件不存在或为空，创建一个空数组并设置max_id为0
        existing_data = []
        max_id = 0
    
    # 读取原始CSV文件并按birth_country分组
    country_groups = defaultdict(list)
    
    with open(input_file, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        all_fields = reader.fieldnames
        
        for row in reader:
            # 处理没有出生国家的记录
            if not row.get('birth_country'):
                continue
            
            # 转换国家名称为中文
            birth_country = convert_country(row.get('birth_country', ''))
            
            # 转换其他国家字段为中文
            country_fields = ['birth_countryNow', 'death_country', 'death_countryNow', 
                             'org_founded_country', 'org_founded_countryNow',
                             'residence_1', 'residence_2']
            
            for field in country_fields:
                if field in row and row[field]:
                    row[field] = convert_country(row[field])
            
            # 保留所有原始字段的信息，但确保birth_country是第一个字段
            laureate_info = {k: v for k, v in row.items() if k != 'birth_country'}
            
            # 将出生国家作为第一个字段
            ordered_info = {'birth_country': birth_country}
            ordered_info.update(laureate_info)  # 添加其他所有字段
            
            # 将信息添加到对应国家组
            country_groups[birth_country].append(ordered_info)
    
    # 将分组数据转换为knowledge.json格式并追加
    new_items = []
    current_id = max_id + 1
    
    for country, laureates in country_groups.items():
        # 将每个获奖者的信息转换为文本
        laureates_count = len(laureates)
        item_text = f"来自{country}的诺贝尔奖获得者共有{laureates_count}位。\n\n"
        
        for idx, laureate in enumerate(laureates, 1):
            # 获取基本信息
            name = laureate.get('name', '未知')
            category = laureate.get('category', '未知')
            year = laureate.get('awardYear', '未知')
            motivation = laureate.get('motivation', '未知')
            
            # 构造文本
            laureate_text = f"获奖者{idx}：\n"
            laureate_text += f"出生国家：{country}\n"
            
            # 添加其他所有非空字段
            for field, value in laureate.items():
                if field != 'birth_country' and value: # 已经添加了出生国家，跳过空值
                    laureate_text += f"{field}：{value}\n"
            
            item_text += laureate_text + "\n"
        
        # 创建新的数据项，所有数据都在text字段中
        new_item = {
            "id": current_id,
            "text": item_text
        }
        new_items.append(new_item)
        current_id += 1
    
    # 将新数据追加到现有数据
    result = existing_data + new_items
    
    # 将结果写回JSON文件
    with open(json_file, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    laureates_total = sum(len(group) for group in country_groups.values())
    print(f"数据已按出生国家分组并追加到{json_file}，共添加了{len(new_items)}个国家组，总计{laureates_total}位获奖者。")

if __name__ == "__main__":
    convert_and_group_data()